package main

import (
	"math/rand"
	"time"
)

type Item struct {
	Name   string
	Weight int
}

// 线性扫描
func weightedRandomSelect(items []Item) Item {
	totalWeight := 0
	// 权重之和
	for _, item := range items {
		totalWeight += item.Weight
	}

	rand.NewSource(time.Now().UnixNano())
	r := rand.Intn(totalWeight)

	for _, item := range items {
		if r < item.Weight {
			return item
		}
		r -= item.Weight
	}

	// This should never happen
	panic("No item selected")
}

// http://strawhatfy.github.io/2015/09/28/two_algorithms_for_drop_probability
// 线性扫描效率太低,可以使用别名记录,用空间换时间
func prepareWeightedRandom4(values []string, weights []float64) func() (string, float64) {
	// 1. 遍历数据,将各个权重的数据归入多个桶中
	// 2. 构造完成alias表
	// 3. 返回闭包,闭包都是查询alias表后再返回
	count := len(weights)
	sum := 0.0
	// 计算总和
	for _, w := range weights {
		sum += w
	}
	avg := sum / float64(count) // 平均权重

	// 记录权重的信息别名
	aliases := make([]int, count)
	probability := make([]float64, count)
	small := make([]int, 0, count)
	large := make([]int, 0, count)
	for i, w := range weights {
		aliases[i] = 1.0
		probability[i] = w
		if w <= avg {
			small = append(small, i)
		} else {
			large = append(large, i)
		}
	}

	// 容器准备完毕,开始填充数据到alias中
	// 权重已经分为了large与small两个部分,所以,根据各自的大小做不同的处理
	for len(small) > 0 && len(large) > 0 {
		less := small[len(small)-1]
		more := large[len(large)-1]

		probability[less] = weights[less] * float64(count)

		aliases[less] = more

		weights[more] = (weights[more] + weights[less]) - avg

		if weights[more] >= 1.0/float64(count) {
			large = append(large, more)
		} else {
			small = append(small, more)
		}

		small = small[:len(small)-1]
		large = large[:len(large)-1]
	}

	return func() (string, float64) {
		n := rand.Float64() * float64(count)
		i := int(n)
		odds := probability[i]
		alias := aliases[i]
		var idx int
		if n-float64(i) > odds {
			idx = alias
		} else {
			idx = i + 1
		}
		return values[idx-1], weights[idx-1]
	}
}

type AliasMethod struct {
	probabilities []float64
	aliases       []int
}

func NewAliasMethod(weights []float64) *AliasMethod {
	n := len(weights)

	// 最终结果会有两个空间,一个存放处理后的概率(用于与随机生成的值相比较),另外一个存放下标
	probabilities := make([]float64, n)
	aliases := make([]int, n)

	sum := 0.0
	for _, weight := range weights {
		sum += weight
	}

	average := sum / float64(n)

	underfull := make([]int, 0)
	overfull := make([]int, 0)

	// 通过与平均值average比较,将weights切分为两个部分--因为要将所有权重缩放到"1"中
	// underfull overfull 仅是记录下标
	for i, weight := range weights {
		probabilities[i] = weight / average
		if probabilities[i] < 1.0 {
			underfull = append(underfull, i)
		} else {
			overfull = append(overfull, i)
		}
	}

	// 遍历处理probabilities (underfull overfull)
	for len(underfull) > 0 && len(overfull) > 0 { // 注意这里是&&
		underIndex := underfull[0]
		overIndex := overfull[0]

		// alias 记录了 小于平均值的权重 对应的 大于平均值权重 的下标
		// probabilities 构造了新的权重与可能性--超过平均值的权重补充到没有达到平均值的权重上

		aliases[underIndex] = overIndex // (此时是赋值, 记录了underIndex与overIndex之间的联系---就是为了把超过的一部分分配到没有超过的空间中)
		probabilities[overIndex] = probabilities[overIndex] - (1.0 - probabilities[underIndex])
		// (x = 1.0 - probabilities[underIndex] 表示需要x才能让probabilities[underIndex]达到平均值)
		// (probabilities[overIndex] - (1.0 - probabilities[underIndex]) 表示从超过平均值的部分中切出一部分用来补充没有超过平均值的权重)

		if probabilities[overIndex] < 1.0 {
			underfull = append(underfull, overIndex)
		} else {
			overfull = append(overfull, overIndex)
		}

		underfull = underfull[1:]
		overfull = overfull[1:]
	}
	// 如果underfull overfull中某一方提前结束了,那么还没有结束的一方的值(也就是权重的下标),
	// 在alias中会直接作为nil,在probabilities中是默认值1

	return &AliasMethod{
		probabilities: probabilities,
		aliases:       aliases,
	}
}

func (am *AliasMethod) Next() int {
	n := len(am.probabilities)
	column := rand.Intn(n)

	if rand.Float64() < am.probabilities[column] {
		return column
	}

	return am.aliases[column]
}

// func main() {
// 	weights := []float64{0.1, 0.2, 0.3, 0.4}
// 	aliasMethod := NewAliasMethod(weights)

// 	counts := make([]int, len(weights))

// 	for i := 0; i < 100000; i++ {
// 		index := aliasMethod.Next()
// 		counts[index]++
// 	}

// 	for i, count := range counts {
// 		probability := float64(count) / 100000
// 		fmt.Printf("Element %d: Probability %.2f%%\n", i, probability*100)
// 	}
// }
